from neo4j import GraphDatabase
from typing import List, Dict, Any
import os

class JobKnowledgeGraph:
    def __init__(self):
        self.uri = os.getenv("NEO4J_URI", "bolt://localhost:7687")
        self.user = os.getenv("NEO4J_USER", "neo4j")
        self.password = os.getenv("NEO4J_PASSWORD", "password")
        self.driver = GraphDatabase.driver(self.uri, 
                                         auth=(self.user, self.password))
    
    def close(self):
        self.driver.close()
    
    def create_constraints(self):
        """创建唯一性约束"""
        with self.driver.session() as session:
            session.write_transaction(
                self._create_constraints
            )
    
    @staticmethod
    def _create_constraints(tx):
        """创建约束的内部方法"""
        tx.run("""
        CREATE CONSTRAINT unique_skill IF NOT EXISTS 
        FOR (s:Skill) REQUIRE s.name IS UNIQUE
        """)
        tx.run("""
        CREATE CONSTRAINT unique_job IF NOT EXISTS 
        FOR (j:Job) REQUIRE j.title IS UNIQUE
        """)
        tx.run("""
        CREATE CONSTRAINT unique_company IF NOT EXISTS 
        FOR (c:Company) REQUIRE c.name IS UNIQUE
        """)
    
    def add_skill(self, skill_name: str):
        """添加技能节点"""
        with self.driver.session() as session:
            session.write_transaction(
                self._add_skill,
                skill_name
            )
    
    @staticmethod
    def _add_skill(tx, skill_name: str):
        """添加技能节点的内部方法"""
        tx.run("""
        MERGE (s:Skill {name: $skill_name})
        """, skill_name=skill_name)
    
    def add_job_posting(self, job_data: Dict[str, Any]):
        """添加岗位节点和关系"""
        with self.driver.session() as session:
            session.write_transaction(
                self._add_job_posting,
                job_data
            )
    
    @staticmethod
    def _add_job_posting(tx, job_data: Dict[str, Any]):
        """添加岗位的内部方法"""
        # 创建岗位节点
        tx.run("""
        MERGE (j:Job {
            title: $title,
            company: $company,
            location: $location
        })
        """, **job_data)
        
        # 创建技能关系
        for skill in job_data.get("skills", []):
            tx.run("""
            MATCH (j:Job {title: $title})
            MERGE (s:Skill {name: $skill})
            MERGE (j)-[r:REQUIRES]->(s)
            """, title=job_data["title"], skill=skill)
        
        # 创建行业关系
        tx.run("""
        MATCH (j:Job {title: $title})
        MERGE (i:Industry {name: $industry})
        MERGE (j)-[r:IN_INDUSTRY]->(i)
        """, title=job_data["title"], industry=job_data.get("industry", "General"))
    
    def find_related_jobs(self, skills: List[str], limit: int = 10):
        """基于技能查找相关岗位"""
        with self.driver.session() as session:
            result = session.read_transaction(
                self._find_related_jobs,
                skills,
                limit
            )
            return [dict(record["job"]) for record in result]
    
    @staticmethod
    def _find_related_jobs(tx, skills: List[str], limit: int):
        """查找相关岗位的内部方法"""
        query = """
        MATCH (j:Job)-[:REQUIRES]->(s:Skill)
        WHERE s.name IN $skills
        WITH j, COUNT(s) AS common_skills
        ORDER BY common_skills DESC
        LIMIT $limit
        RETURN j
        """
        return tx.run(query, skills=skills, limit=limit)
    
    def get_skill_graph(self, skill_name: str, depth: int = 2):
        """获取技能关联图谱"""
        with self.driver.session() as session:
            result = session.read_transaction(
                self._get_skill_graph,
                skill_name,
                depth
            )
            return [dict(record) for record in result]
    
    @staticmethod
    def _get_skill_graph(tx, skill_name: str, depth: int):
        """获取技能图谱的内部方法"""
        query = """
        MATCH path = (s:Skill {name: $skill_name})-[*1..$depth]-(related)
        RETURN nodes(path) AS nodes, relationships(path) AS relationships
        """
        return tx.run(query, skill_name=skill_name, depth=depth)

if __name__ == "__main__":
    kg = JobKnowledgeGraph()
    try:
        kg.create_constraints()
        
        # 示例数据
        job_data = {
            "title": "Data Scientist",
            "company": "TechCorp",
            "location": "San Francisco",
            "skills": ["Python", "Machine Learning", "Statistics"],
            "industry": "IT"
        }
        kg.add_job_posting(job_data)
        
        # 查询示例
        print("相关岗位:", kg.find_related_jobs(["Python", "Machine Learning"]))
    finally:
        kg.close()